from keras.models import *
from attention import *
from bilinear_upsampling import BilinearUpsampling
import tensorflow as tf
from keras.layers import Input, Conv2D, MaxPooling2D, BatchNormalization, Activation, Add, Concatenate, Dropout
from keras.regularizers import l2

class BatchNorm(BatchNormalization):
    def call(self, inputs, training=None):
        return super(self.__class__, self).call(inputs, training=True)

class Copy(Layer):
    def call(self, inputs, **kwargs):
        copy = tf.identity(inputs)
        return copy

    def compute_output_shape(self, input_shape):
        return input_shape

class layertile(Layer):
    def call(self, inputs, **kwargs):
        image = tf.reduce_mean(inputs, axis=-1)
        image = tf.expand_dims(image, -1)
        image = tf.tile(image, [1, 1, 1, 32])
        return image

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)[:-1] + [32]
        return tuple(output_shape)


def BN(input_tensor, block_id):
    bn = BatchNorm(name=block_id + '_BN')(input_tensor)
    a = Activation('relu', name=block_id + '_relu')(bn)
    return a


def AtrousBlock(input_tensor, filters, rate, block_id, stride=1):
    x = Conv2D(filters, (3, 3), strides=(stride, stride), dilation_rate=(rate, rate),
               padding='same', use_bias=False, name=block_id + '_dilation')(input_tensor)
    return x


def CFE(input_tensor, filters, block_id):
    rate = [3, 5, 7]
    cfe0 = Conv2D(filters, (1, 1), padding='same', use_bias=False, name=block_id + '_cfe0')(
        input_tensor)
    cfe1 = AtrousBlock(input_tensor, filters, rate[0], block_id + '_cfe1')
    cfe2 = AtrousBlock(input_tensor, filters, rate[1], block_id + '_cfe2')
    cfe3 = AtrousBlock(input_tensor, filters, rate[2], block_id + '_cfe3')
    cfe_concat = Concatenate(name=block_id + 'concatcfe', axis=-1)([cfe0, cfe1, cfe2, cfe3])
    cfe_concat = BN(cfe_concat, block_id)
    return cfe_concat


# ResNet-18 基本残差块
def basic_block(input_tensor, filters, strides=1, downsample=None, block_id=''):
    x = Conv2D(filters, kernel_size=3, strides=strides, padding='same',
               kernel_initializer='he_normal', kernel_regularizer=l2(1e-4),
               name=block_id + '_conv1')(input_tensor)
    x = BatchNorm(name=block_id + '_bn1')(x)
    x = Activation('relu', name=block_id + '_relu1')(x)

    x = Conv2D(filters, kernel_size=3, strides=1, padding='same',
               kernel_initializer='he_normal', kernel_regularizer=l2(1e-4),
               name=block_id + '_conv2')(x)
    x = BatchNorm(name=block_id + '_bn2')(x)

    if downsample is not None:
        shortcut = downsample(input_tensor)
    else:
        shortcut = input_tensor

    x = Add(name=block_id + '_add')([x, shortcut])
    x = Activation('relu', name=block_id + '_relu2')(x)
    return x


# ResNet-18 构建函数
def ResNet18(img_input):
    # 初始卷积层
    x = Conv2D(64, kernel_size=7, strides=2, padding='same',
               kernel_initializer='he_normal', kernel_regularizer=l2(1e-4),
               name='conv1')(img_input)
    x = BatchNorm(name='bn1')(x)
    x = Activation('relu', name='relu1')(x)
    x = MaxPooling2D(pool_size=3, strides=2, padding='same', name='maxpool')(x)

    # 第 1 个残差块组
    num_filters = 64
    x = basic_block(x, num_filters, block_id='layer1_0')
    C1 = x
    x = basic_block(x, num_filters, block_id='layer1_1')

    # 第 2 个残差块组
    num_filters *= 2
    downsample = Conv2D(num_filters, kernel_size=1, strides=2,
                        kernel_initializer='he_normal', kernel_regularizer=l2(1e-4),
                        name='layer2_downsample')
    x = basic_block(x, num_filters, strides=2, downsample=downsample, block_id='layer2_0')
    C2 = x
    x = basic_block(x, num_filters, block_id='layer2_1')

    # 第 3 个残差块组
    num_filters *= 2
    downsample = Conv2D(num_filters, kernel_size=1, strides=2,
                        kernel_initializer='he_normal', kernel_regularizer=l2(1e-4),
                        name='layer3_downsample')
    x = basic_block(x, num_filters, strides=2, downsample=downsample, block_id='layer3_0')
    C3 = x
    x = basic_block(x, num_filters, block_id='layer3_1')

    # 第 4 个残差块组
    num_filters *= 2
    downsample = Conv2D(num_filters, kernel_size=1, strides=2,
                        kernel_initializer='he_normal', kernel_regularizer=l2(1e-4),
                        name='layer4_downsample')
    x = basic_block(x, num_filters, strides=2, downsample=downsample, block_id='layer4_0')
    C4 = x
    x = basic_block(x, num_filters, block_id='layer4_1')
    C5 = x

    return C1, C2, C3, C4, C5


def ResNet18BasedModel(img_input, dropout=False, with_CPFE=False, with_CA=False, with_SA=False, droup_rate=0.3):
    C1, C2, C3, C4, C5 = ResNet18(img_input)

    C1 = Conv2D(64, (3, 3), padding='same', name='C1_conv')(C1)
    C1 = BN(C1, 'C1_BN')
    C2 = Conv2D(64, (3, 3), padding='same', name='C2_conv')(C2)
    C2 = BN(C2, 'C2_BN')

    if with_CPFE:
        C3_cfe = CFE(C3, 32, 'C3_cfe')
        C4_cfe = CFE(C4, 32, 'C4_cfe')
        C5_cfe = CFE(C5, 32, 'C5_cfe')

        # 对 C4_cfe 上采样 2 倍，使其与 C3_cfe 对齐 (from 8x8 -> 16x16)
        C4_cfe = BilinearUpsampling(upsampling=(2, 2), name='C4_cfe_up2')(C4_cfe)

        # 对 C5_cfe 也上采样 2 倍，使其也达到 16x16
        C5_cfe = BilinearUpsampling(upsampling=(2, 2), name='C5_cfe_up2')(C5_cfe)

        C345 = Concatenate(name='C345_aspp_concat', axis=-1)([C3_cfe, C4_cfe, C5_cfe])

        if with_CA:
            C345 = ChannelWiseAttention(C345, name='C345_ChannelWiseAttention_withcpfe')
    else:
        C345 = C5  # fallback

    C345 = Conv2D(64, (1, 1), padding='same', name='C345_conv')(C345)
    C345 = BN(C345, 'C345')
    C345 = BilinearUpsampling(upsampling=(4, 4), name='C345_up4')(C345)

    if with_SA:
        SA = SpatialAttention(C345, 'spatial_attention')
        C2 = BilinearUpsampling(upsampling=(2, 2), name='C2_up2')(C2)
        C12 = Concatenate(name='C12_concat', axis=-1)([C1, C2])
        C12 = Conv2D(64, (3, 3), padding='same', name='C12_conv')(C12)
        C12 = BN(C12, 'C12')
        C12 = Multiply(name='C12_atten_mutiply')([SA, C12])
    else:
        C2 = BilinearUpsampling(upsampling=(2, 2), name='C2_up2')(C2)
        C12 = Concatenate(name='C12_concat', axis=-1)([C1, C2])
        C12 = Conv2D(64, (3, 3), padding='same', name='C12_conv')(C12)
        C12 = BN(C12, 'C12')

    fea = Concatenate(name='fuse_concat', axis=-1)([C12, C345])
    sa = Conv2D(1, (3, 3), padding='same', name='sa')(fea)

    model = Model(inputs=img_input, outputs=sa, name="ResNet18BasedModel")
    return model
